import math
import os
from datetime import datetime, timedelta

import numpy as np

from RadReader import RadReader


def select_median0(arr):
    # 排除 NaN 值
    arr = arr[~np.isnan(arr)]
    # 获取元素个数
    n = len(arr)
    if n == 0:
        return None
    # 如果数组元素个数为奇数，直接返回中位数
    if n % 2 == 1:
        return np.median(arr)
    # 数组元素个数为偶数时
    else:
        # 计算数组的平均值
        avg = np.mean(arr)
        # 计算中位数的两个值
        sorted_arr = np.sort(arr)
        lower_mid = sorted_arr[(n // 2) - 1]
        upper_mid = sorted_arr[n // 2]
        # 返回更接近平均值的那个中位数
        if abs(lower_mid - avg) <= abs(upper_mid - avg):
            return lower_mid
        else:
            return upper_mid


def select_median(arr, get_value=None):
    if get_value:
        # 用 get_value 取得数组的值，并保留原始元素和对应的值
        arr_with_values = [(x, get_value(x)) for x in arr]
        # 排除 NaN 值
        arr_with_values = [x for x in arr_with_values if not np.isnan(x[1])]
        # 提取非 NaN 的值
        values = np.array([x[1] for x in arr_with_values])
    else:
        # 直接排除 NaN 并处理
        arr_with_values = arr[~np.isnan(arr)]
        values = arr_with_values
    # 获取元素个数
    n = len(values)
    if n == 0:
        return None
    # 如果数组元素个数为奇数，直接返回中位数对应的原始数据元素
    if n % 2 == 1:
        median_value = np.median(values)
        return [x[0] for x in arr_with_values if x[1] == median_value][0]
    # 数组元素个数为偶数时
    else:
        # 计算数组的平均值
        avg = np.mean(values)
        # 计算中位数的两个值
        sorted_values = np.sort(values)
        lower_mid_value = sorted_values[(n // 2) - 1]
        upper_mid_value = sorted_values[n // 2]
        # 返回更接近平均值的那个中位数对应的原始数据元素
        lower_mid_element = [x[0] for x in arr_with_values if x[1] == lower_mid_value][0]
        upper_mid_element = [x[0] for x in arr_with_values if x[1] == upper_mid_value][0]
        if abs(lower_mid_value - avg) <= abs(upper_mid_value - avg):
            return lower_mid_element
        else:
            return upper_mid_element


def format_value(value, pattern, length):
    if np.isnan(value):
        return "/" * length
    else:
        return pattern.format(value)


class WindProfile:

    def __init__(self, filename):
        self.filename = filename
        self.time_str = filename.split("_")[4]
        self.header_line1 = None
        self.header_line2 = None
        self.theta = 14.0
        self.RNESW = 'RNESW/'
        self.azimuth_correction_angle = 173.8

        self.R = []
        self.N = []
        self.E = []
        self.S = []
        self.W = []

    def sort_RNESW(self):
        self.all = []
        for ch in self.RNESW:
            if ch == 'R':
                self.all.append(self.R)
            if ch == 'N':
                self.all.append(self.N)
            if ch == 'E':
                self.all.append(self.E)
            if ch == 'S':
                self.all.append(self.S)
            if ch == 'W':
                self.all.append(self.W)

    def start_rad(self, line1, line2):
        print(f"开始解析文件: {line1}, {line2}")
        self.header_line1 = line1
        self.header_line2 = line2

    def start_section(self, section_index, line1, line2):
        print(f"开始解析第 {section_index + 1} 段: {line1}, {line2}")
        a = line1.split()
        self.theta = float(a[2])
        a = line2.split()
        self.RNESW = a[8]
        self.sort_RNESW()
        self.azimuth_correction_angle = float(a[9])

    def start_items(self, section_index, items_index, line):
        print(f"开始解析第 {items_index + 1} 数据块: {line}")

    def item(self, section_index, items_index, line):
        # print(f"数据行: {line}")
        a = [float('nan') if value[0] == '/' else float(value) for value in line.split()]
        item_data = self.all[items_index]
        n = next((index for index, value in enumerate(item_data) if value[0] == a[0]), -1)
        if n >= 0:
            if not math.isnan(a[3]):
                item_data[n] = a
        else:
            item_data.append(a)

    def end_items(self, section_index, items_index):
        print(f"结束解析第 {items_index + 1} 数据块")

    def end_section(self, section_index):
        print(f"结束解析第 {section_index + 1} 段")

    def end_rad(self):
        print("结束解析文件")
        self.R = np.array(self.R, dtype=np.float32)
        self.N = np.array(self.N, dtype=np.float32)
        self.E = np.array(self.E, dtype=np.float32)
        self.S = np.array(self.S, dtype=np.float32)
        self.W = np.array(self.W, dtype=np.float32)
        self.sort_RNESW()

    def beam_consistency_filter(self):
        self.V_R = np.nanmean(self.R[:, 3])
        self.V_N = np.nanmean(self.N[:, 3])
        self.V_E = np.nanmean(self.E[:, 3])
        self.V_S = np.nanmean(self.S[:, 3])
        self.V_W = np.nanmean(self.W[:, 3])

        self.S_R = np.sqrt(np.nanmean((self.R[:, 3] - self.V_R) ** 2))
        self.S_N = np.sqrt(np.nanmean((self.N[:, 3] - self.V_N) ** 2))
        self.S_E = np.sqrt(np.nanmean((self.E[:, 3] - self.V_E) ** 2))
        self.S_S = np.sqrt(np.nanmean((self.S[:, 3] - self.V_S) ** 2))
        self.S_W = np.sqrt(np.nanmean((self.W[:, 3] - self.V_W) ** 2))

        indices = np.where((self.R[:, 3] - self.V_R) > self.S_R)[0]
        self.R[indices, 1:4] = np.nan
        indices = np.where((self.N[:, 3] - self.V_N) > self.S_N)[0]
        self.N[indices, 1:4] = np.nan
        indices = np.where((self.E[:, 3] - self.V_E) > self.S_E)[0]
        self.E[indices, 1:4] = np.nan
        indices = np.where((self.S[:, 3] - self.V_S) > self.S_S)[0]
        self.S[indices, 1:4] = np.nan
        indices = np.where((self.W[:, 3] - self.V_W) > self.S_W)[0]
        self.W[indices, 1:4] = np.nan
        pass

    def symmetry_filter(self):
        self.u_E = ((self.E[:, 3] + self.W[:, 3]) / np.sin(self.theta)).astype(np.float32)
        self.u_W = (2 * self.R[:, 3] / np.tan(self.theta)).astype(np.float32)
        self.delta_u = self.u_E - self.u_W

        self.v_N = ((self.N[:, 3] + self.S[:, 3]) / np.sin(self.theta)).astype(np.float32)
        self.v_S = self.u_W
        self.delta_v = self.v_N - self.v_S

        # 当 δu >= 1.0 且 (E[:, 3] - V_E) >= S_E，设置 E 的元素为 NaN
        indices = np.where((self.delta_u >= 1.0) & ((self.E[:, 3] - self.V_E) >= self.S_E))[0]
        self.E[indices, 1:4] = np.nan
        # 当 δu >= 1.0 且 (W[:, 3] - V_W) >= S_W，设置 W 的元素为 NaN
        indices = np.where((self.delta_u >= 1.0) & ((self.W[:, 3] - self.V_W) >= self.S_W))[0]
        self.W[indices, 1:4] = np.nan
        # 当 δv >= 1.0 且 (N[:, 3] - V_N) >= S_N，设置 N 的元素为 NaN
        indices = np.where((self.delta_v >= 1.0) & ((self.N[:, 3] - self.V_N) >= self.S_N))[0]
        self.N[indices, 1:4] = np.nan
        # 当 δv >= 1.0 且 (S[:, 3] - V_S) >= S_S，设置 S 的元素为 NaN
        indices = np.where((self.delta_v >= 1.0) & ((self.S[:, 3] - self.V_S) >= self.S_S))[0]
        self.S[indices, 1:4] = np.nan
        pass

    def radial_velocity_revise(self):
        self.W_N = (self.N[:, 3] / np.cos(self.theta)).astype(np.float32)
        self.W_E = (self.E[:, 3] / np.cos(self.theta)).astype(np.float32)
        self.W_S = (self.S[:, 3] / np.cos(self.theta)).astype(np.float32)
        self.W_W = (self.W[:, 3] / np.cos(self.theta)).astype(np.float32)

        R1 = np.full(len(self.R), np.nan, dtype=np.float32)
        for n in range(len(self.R)):
            arr = np.array([self.W_N[n], self.W_E[n], self.R[n, 3], self.W_S[n], self.W_W[n]], dtype=np.float32)
            # 取中位值
            median_value = select_median0(arr)
            if median_value is not None:
                R1[n] = median_value

        # 保留原始的 self.R 以备后用
        self.R0 = np.copy(self.R)

        # 将 self.R 的第四列替换为 R1
        self.R[:, 3] = R1
        pass

    def fill_in_nan(self, history_wind_profiles):
        # 列表中最后一个元素是当前结果
        count = len(history_wind_profiles) - 1
        for n, one in enumerate(self.all):
            indices = np.where(np.isnan(one[:, 3]))[0]
            arr = np.full((count, 5, 4), np.nan, dtype=np.float32)
            for i in indices:
                for j in range(count):
                    history_wind_profile = history_wind_profiles[j]
                    arr[j, 0] = history_wind_profile.all[n][i]
                    if i - 2 >= 0:
                        arr[j, 1] = history_wind_profile.all[n][i - 2]
                    if i - 1 >= 0:
                        arr[j, 2] = history_wind_profile.all[n][i - 1]
                    if i + 1 < len(one):
                        arr[j, 3] = history_wind_profile.all[n][i + 1]
                    if i + 2 < len(one):
                        arr[j, 4] = history_wind_profile.all[n][i + 2]

                median_values = []
                get_value = lambda x: x[3]
                a_2 = select_median(arr[:, 1, :], get_value)
                if a_2 is not None:
                    median_values.append(a_2)
                a_1 = select_median(arr[:, 2, :], get_value)
                if a_1 is not None:
                    median_values.append(a_1)
                a0 = select_median(arr[:, 0, :], get_value)
                if a0 is not None:
                    median_values.append(a0)
                a1 = select_median(arr[:, 3, :], get_value)
                if a1 is not None:
                    median_values.append(a1)
                a2 = select_median(arr[:, 4, :], get_value)
                if a2 is not None:
                    median_values.append(a2)
                a = select_median(median_values, get_value)
                if a is not None:
                    one[i] = np.copy(a)
        pass

    def calc_uv(self):
        r_indices = np.where(~np.isnan(self.R[:, 3]))[0]

        self.U = np.full(len(self.R), np.nan, dtype=np.float32)
        ew_indices = np.where(~np.isnan(self.E[r_indices, 3]) & ~np.isnan(self.W[r_indices, 3]))[0]
        self.U[ew_indices] = 0.5 * (
                (self.E[ew_indices, 3] - self.R[ew_indices, 3] * np.cos(self.theta)) / np.sin(self.theta) +
                (self.W[ew_indices, 3] - self.R[ew_indices, 3] * np.cos(self.theta)) / -np.sin(self.theta))
        e_indices = np.where(~np.isnan(self.E[r_indices, 3]) & np.isnan(self.W[r_indices, 3]))[0]
        self.U[e_indices] = (self.E[e_indices, 3] - self.R[e_indices, 3] * np.cos(self.theta)) / np.sin(self.theta)
        w_indices = np.where(np.isnan(self.E[r_indices, 3]) & ~np.isnan(self.W[r_indices, 3]))[0]
        self.U[w_indices] = (self.W[w_indices, 3] - self.R[w_indices, 3] * np.cos(self.theta)) / -np.sin(self.theta)

        self.V = np.full(len(self.R), np.nan, dtype=np.float32)
        ns_indices = np.where(~np.isnan(self.N[r_indices, 3]) & ~np.isnan(self.S[r_indices, 3]))[0]
        self.V[ns_indices] = 0.5 * (
                (self.N[ns_indices, 3] - self.R[ns_indices, 3] * np.cos(self.theta)) / np.sin(self.theta) +
                (self.S[ns_indices, 3] - self.R[ns_indices, 3] * np.cos(self.theta)) / -np.sin(self.theta))
        n_indices = np.where(~np.isnan(self.N[r_indices, 3]) & np.isnan(self.S[r_indices, 3]))[0]
        self.V[n_indices] = (self.N[n_indices, 3] - self.R[n_indices, 3] * np.cos(self.theta)) / np.sin(self.theta)
        s_indices = np.where(np.isnan(self.N[r_indices, 3]) & ~np.isnan(self.S[r_indices, 3]))[0]
        self.V[s_indices] = (self.S[s_indices, 3] - self.R[s_indices, 3] * np.cos(self.theta)) / -np.sin(self.theta)
        pass

    def calc_wind(self):
        uv_indices = np.where(~np.isnan(self.U) & ~np.isnan(self.V))

        self.W_spd = np.full(len(self.R), np.nan, dtype=np.float32)
        self.W_spd[uv_indices] = np.sqrt(self.U[uv_indices] ** 2 + self.V[uv_indices] ** 2)

        self.W_dir = np.full(len(self.R), np.nan, dtype=np.float32)
        # 忽略 RuntimeWarning: divide by zero encountered in divide，np.arctan可以计算出正确结果pi/2或-pi/2
        with np.errstate(divide='ignore', invalid='ignore'):
            self.W_dir[uv_indices] = np.arctan(self.U[uv_indices] / self.V[uv_indices]) + self.azimuth_correction_angle
        pass

    def join_robs_data(self):
        self.ROBS = np.full((len(self.R), 7), np.nan, dtype=np.float32)
        self.ROBS[:, 0] = self.R[:, 0]
        self.ROBS[:, 1] = self.W_spd
        self.ROBS[:, 2] = self.W_dir
        self.ROBS[:, 3] = self.R[:, 3]
        self.ROBS[:, 4] = 100
        self.ROBS[:, 5] = 100
        self.ROBS[:, 6] = self.R[:, 2]
        pass

    @classmethod
    def avg_robs_data(cls, wind_profiles):
        sample = wind_profiles[-1]
        result = WindProfile(sample.filename)
        result.header_line1 = sample.header_line1
        result.header_line2 = sample.header_line2

        result.ROBS = np.full((len(sample.R), 7), np.nan, dtype=np.float32)
        ROBS_cube = np.array([wind_profile.ROBS for wind_profile in wind_profiles])
        result.ROBS[:, 0] = sample.ROBS[:, 0]
        result.ROBS[:, 1] = np.nanmean(ROBS_cube[:, :, 1], axis=0)
        result.ROBS[:, 2] = np.nanmean(ROBS_cube[:, :, 2], axis=0)
        result.ROBS[:, 3] = np.nanmean(ROBS_cube[:, :, 3], axis=0)
        result.ROBS[:, 4] = 100
        result.ROBS[:, 5] = 100
        result.ROBS[:, 6] = np.nanmean(ROBS_cube[:, :, 6], axis=0)
        return result

    def write_robs_file(self, dir):
        a = self.filename.split("_")
        a[4] = self.time_str
        a[5] = 'P'
        a[8] = 'ROBS.TXT'
        robs_filename = "_".join(a)
        output_file = os.path.join(dir, robs_filename)
        with open(output_file, 'w') as f:
            line = self.header_line1.replace("WNDRAD", "WNDROBS")
            f.write(line + "\n")
            line = self.header_line2 + " " + self.time_str
            f.write(line + "\n")
            f.write("ROBS\n")
            for item in self.ROBS:
                formatted_cols = [
                    format_value(int(item[0]), '{:05d}', 5),
                    format_value(item[1], '{:05.1f}', 5),
                    format_value(item[2], '{:05.1f}', 5),
                    format_value(item[3], '{:06.1f}', 5),
                    format_value(int(item[4]), '{:03d}', 3),
                    format_value(int(item[5]), '{:03d}', 3),
                    format_value(item[6], '{:08.1e}', 8)
                ]
                line = ' '.join(formatted_cols)
                f.write(line + "\n")
            f.write("NNNN\n")
            return output_file


def scan_files(directory, target_time, lead_minutes):
    result = []
    for filename in os.listdir(directory):
        if filename.startswith("Z_RADA_I_") and filename.endswith("_O_WPRD_LC_RAD.TXT"):
            try:
                timestamp_str = filename.split("_")[4]
                timestamp = datetime.strptime(timestamp_str, "%Y%m%d%H%M%S")
                if target_time - timedelta(minutes=lead_minutes) <= timestamp <= target_time:
                    result.append(filename)
            except ValueError:
                continue
    return result


def scan_profiles(profiles, target_time, lead_minutes):
    result = []
    for profile in profiles:
        timestamp = datetime.strptime(profile.time_str, "%Y%m%d%H%M%S")
        if target_time - timedelta(minutes=lead_minutes) <= timestamp <= target_time:
            result.append(profile)
    return result


def calc_robs(data_dir, time_str, lead_minutes=30):
    # 1 获取指定时间前30 * 2分钟内的数据文件列表
    target_time = datetime.strptime(time_str, "%Y%m%d%H%M%S")
    filenames = scan_files("data", target_time, lead_minutes * 2)

    # 形成结果集
    all_wind_profiles = []
    for filename in filenames:
        filepath = os.path.join(data_dir, filename)
        # 2 对每个数据文件进行段合并，对同一高度，以有效>无效，后>前的关系进行覆盖
        reader = RadReader(filepath)
        wind_profile = reader.parse(WindProfile(filename))
        all_wind_profiles.append(wind_profile)

    # 数据过滤和订正
    for wind_profile in all_wind_profiles:
        # 3 对每个结果进行5个方向的波束一致性检查，将偏差数据置为无效
        wind_profile.beam_consistency_filter()
        # 4 对每个结果中的EWSN方向进行对称性检查，将偏差数据置为无效
        wind_profile.symmetry_filter()
        # 5 对每个结果中垂直方位的径向速度进行订正，取R方向及EWSN4方向倾斜波束分解值，再求中位数
        wind_profile.radial_velocity_revise()

    # 6 利用过去时间结果对当前结果进行缺测补全
    current_wind_profiles = scan_profiles(all_wind_profiles, target_time, lead_minutes)
    for wind_profile in current_wind_profiles:
        timestamp = datetime.strptime(wind_profile.time_str, "%Y%m%d%H%M%S")
        history_wind_profiles = scan_profiles(all_wind_profiles, timestamp, lead_minutes)
        wind_profile.fill_in_nan(history_wind_profiles)
        # 7.1 计算u、v分量
        wind_profile.calc_uv()
        # 7.2 计算风向风速
        wind_profile.calc_wind()
        # 7.3 组装robs数据
        wind_profile.join_robs_data()

    # 8 多个robs结果取算术平均，生成最终robs结果
    final_wind_profile = WindProfile.avg_robs_data(current_wind_profiles)
    final_wind_profile.time_str = time_str

    robs_file = final_wind_profile.write_robs_file(data_dir)
    return robs_file
